package net.xo68.coordinator.core.gid;

import com.alipay.sofa.jraft.Closure;
import com.alipay.sofa.jraft.Iterator;
import com.alipay.sofa.jraft.Status;
import com.alipay.sofa.jraft.core.StateMachineAdapter;
import com.alipay.sofa.jraft.error.RaftError;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotReader;
import com.alipay.sofa.jraft.storage.snapshot.SnapshotWriter;
import com.alipay.sofa.jraft.util.Utils;
import lombok.extern.slf4j.Slf4j;
import net.xo68.coordinator.core.rpc.ObjectSerialize;
import net.xo68.coordinator.core.rpc.ObjectSerializeFactory;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 全局编号的状态机
 * @author wuxie
 * @version 2020-9-28
 */
@Slf4j
public class GidStateMachine extends StateMachineAdapter {

    private final ConcurrentHashMap<Long, AtomicLong> gidMap=new ConcurrentHashMap<>();

    private final ReentrantLock lock=new ReentrantLock();

    /**
     * leader term
     */
    private final AtomicLong leaderTerm = new AtomicLong(-1);

    public boolean isLeader() {
        return this.leaderTerm.get() > 0;
    }

    @Override
    public void onApply(Iterator iterator) {
        //遍历日志
        while (iterator.hasNext()){
            long idKey=0;

            NewIdClosure closure=null;
            // done 回调不为null，必须在应用日志后调用，如果不为 null，说明当前是leader。
            if(iterator.done()!=null){
                closure=(NewIdClosure) iterator.done();
                idKey= closure.getRequest().getIdKey();
            }else {
                // 其他节点应用此日志，需要反序列化 IncrementAndGetRequest，获取 delta
                ByteBuffer data = iterator.getData();

                ObjectSerialize objectSerialize = ObjectSerializeFactory.get();
                NewIdRequest request=objectSerialize.deserialize(data.array(), NewIdRequest.class);
                idKey = request.getIdKey();
            }

            AtomicLong gid=getGid(idKey);
            //更新后的值
            long updated=gid.incrementAndGet();

            if(closure!=null){
                closure.getResponse().setIdKey(idKey);
                closure.getResponse().setValue(updated);
                closure.getResponse().setStatus(true);
                closure.run(Status.OK());
            }

            iterator.next();
        }
    }

    private AtomicLong getGid(long idKey){
        AtomicLong ret = this.gidMap.get(idKey);
        if (ret == null) {
            lock.lock();
            if (ret == null){
                ret = new AtomicLong(0);
                final AtomicLong old = this.gidMap.putIfAbsent(idKey, ret);
                if (old != null) {
                    ret = old;
                }
            }
            lock.unlock();

        }
        return ret;
    }

    /**
     * 重置 gidMap
     * @param gidData
     */
    private void resetGid(GidSnapshotStore.GidData gidData){
        if(gidData==null || gidData.getData()==null || gidData.getData().isEmpty()){
            return;
        }
        for (Map.Entry<Long,Long> entry: gidData.getData().entrySet()){
            AtomicLong gid=getGid(entry.getKey());
            gid.set(entry.getValue());
        }
    }

    @Override
    public void onLeaderStart(final long term) {
        this.leaderTerm.set(term);
        super.onLeaderStart(term);

    }

    @Override
    public void onLeaderStop(final Status status) {
        this.leaderTerm.set(-1);
        super.onLeaderStop(status);
    }

    @Override
    public boolean onSnapshotLoad(SnapshotReader reader) {
        // leader不用从 snapshot 加载，他不会接受 snapshot 安装请求
        if (isLeader()) {
            log.warn("Leader is not supposed to load snapshot");
            return false;
        }
        // 未找到数据文件，忽略
        if (reader.getFileMeta("data") == null) {
            log.error("Fail to find data file in {}", reader.getPath());
            return false;
        }
        // 将 snapshot 保存在 reader.getPath()/data 文件里
        GidSnapshotStore snapshot = new GidSnapshotStore(reader.getPath() + File.separator + "data");
        try {
            this.resetGid(snapshot.load());
            return true;
        } catch (IOException e) {
            log.error("Fail to load snapshot from {}", snapshot.getPath());
            return false;
        }
    }

    @Override
    public void onSnapshotSave(SnapshotWriter writer, Closure done) {
        // 获取此刻状态机状态
        final GidSnapshotStore.GidData gidData=new GidSnapshotStore.GidData();
        for (Map.Entry<Long,AtomicLong> entry: this.gidMap.entrySet()){
            gidData.getData().put(entry.getKey(),entry.getValue().get());
        }
        // 异步保存，避免阻塞状态机
        Utils.runInThread(new Runnable() {

            @Override
            public void run() {
                GidSnapshotStore snapshot = new GidSnapshotStore(writer.getPath() + File.separator + "data");
                if (snapshot.save(gidData)) {
                    if (writer.addFile("data")) {
                        done.run(Status.OK());
                    } else {
                        done.run(new Status(RaftError.EIO, "Fail to add file to writer"));
                    }
                } else {
                    done.run(new Status(RaftError.EIO, "Fail to save counter snapshot %s", snapshot.getPath()));
                }
            }
        });
    }
}
